import os
from setuptools import setup, find_packages
import subprocess
import sys


def install_faiss_and_transformers():
    """根据环境安装 faiss 和 sentence_transformers"""
    try:
        # 检查是否安装了 NVIDIA CUDA 工具包
        output = subprocess.check_output("nvcc --version", shell=True)
        if "release" in output.decode():
            print("检测到 NVIDIA CUDA 环境，安装 faiss-gpu 和 sentence_transformers")
            subprocess.check_call([sys.executable, "-m", "pip", "install", "faiss-gpu", "sentence_transformers"])
        else:
            raise Exception("未检测到 NVIDIA CUDA 环境")
    except Exception as e:
        print(f"未检测到 GPU 或 CUDA 环境，仅安装 faiss-cpu 和 sentence_transformers")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "faiss-cpu", "sentence_transformers"])


def get_install_requires():
    """获取项目的依赖"""
    requires = [
        'numpy==1.23.5',
        'langchain==0.2.5',
        'langchain-community==0.2.5',
        'einops==0.8.0',
        'requests==2.32.3',
        'tiktoken==0.7.0',
        'tinydb==4.8.0',
        'tqdm==4.66.4',
    ]

    # 检查是否需要安装 transformers 和其他本地化部署依赖
    deployment_type = os.environ.get("DEPLOYMENT_TYPE", "api")  # 默认为 API 部署
    if deployment_type == "local":
        requires.append('transformers==4.35.1')
        install_faiss_and_transformers()  # 调用安装 faiss 和 sentence_transformers 的函数

    return requires


setup(
    name="survey_pass",
    version="1.0.0",
    description="This is a magical tool for assisting in writing papers",
    long_description="This is a tool for assisting in writing papers, which can automatically outline, write, and evaluate.",
    author="BreakerLM-Lab",
    author_email="breakerlm@gmail.com",
    url="https://github.com/BreakerLM-Lab/BreakerLM",
    packages=find_packages(),
    install_requires=get_install_requires(),
)